package com.company.nuwa.common.handler;

import com.company.nuwa.common.constants.Headers;
import com.company.nuwa.common.utils.RequestContextUtil;
import com.company.nuwa.common.utils.ZkPlatformUtils;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;
import org.springframework.core.task.TaskDecorator;

import java.util.Objects;

/**
 * <p>线程任务装饰器</p>
 *
 * @author: Spector
 * @version: 1.0
 * @date 2023/12/26 18:05
 */
@Slf4j
public class ThreadTaskDecorator implements TaskDecorator {

    @Override
    public Runnable decorate(Runnable runnable) {
        var mdcContext = MDC.getCopyOfContextMap();
        var platformId = ZkPlatformUtils.getPlatformId();
        String platformIdStr = RequestContextUtil.getHeader(Headers.PLATFORM_ID);
        log.info("[异步线程启动] traceId: {}, zkUtilPlatformId: {}, headerPlatformId: {}",
                Objects.nonNull(mdcContext) ? mdcContext.get("traceId") : null, platformId, platformIdStr);
        return () -> {
            try {
                // 传递上下文
                MDC.setContextMap(mdcContext);
                ZkPlatformUtils.setPlatformId(platformId);

                runnable.run();
            } finally {
                // 调用完毕后清理
                MDC.clear();
                ZkPlatformUtils.removePlatformId();
            }
        };
    }

}
